import torch
import math

class WarmupCosineLR(object):
    def __init__(self, optimizer, base_lr, warm_lr, max_iters, warm_iters, eta_min):
        self.lr = 0
        self.base_lr = base_lr
        self.warm_lr = warm_lr
        self.max_iters = max_iters
        self.warm_iters = warm_iters
        self.eta_min = eta_min
        self.optimizer = optimizer
        

    def step(self, iter):
        if self.warm_iters>iter:
            lr = self.warm_lr+(self.base_lr-self.warm_lr) * iter / self.warm_iters
        else:
            if iter>self.max_iters:
                lr = self.eta_min
            else:
                lr = self.eta_min + (self.base_lr-self.eta_min)*\
                    (1+math.cos(math.pi*(iter-self.warm_iters)/(self.max_iters-self.warm_iters)))/2

        for param_group in self.optimizer.param_groups:
                param_group['lr'] = lr
